'''
"doc_scanner.py" 文件中的 DocScanner 类提供了文档扫描的核心功能，功能包括
    - 加载图片
    - 找到文档的角落
    - 裁剪图片

- 通过将文档扫描功能封装在一个独立的类中，我们可以将其作为一个模块，方便在其他项目中重用。
- 这种模块化的设计使得代码更加模块化、可扩展和可测试。
- 它可以独立于其他部分进行开发、测试和维护，而不会对其他组件产生影响。
- 此外，将文档扫描功能单独放在一个文件中，使得代码结构更清晰、易于理解和维护。
- 通过将不同的功能分割到不同的文件中，我们可以更好地组织代码，减少文件的复杂性，并促进团队合作开发。
'''
import cv2 as cv  # 导入 OpenCV 库，用于图片处理和计算机视觉任务
import numpy as np  # 导入 numpy 库，用于进行数值计算

class DocScanner:
    def __init__(self):
        pass  # DocScanner 的初始化方法，目前为空

    def load_image(self, file_path):
        """加载图片并找到文档的四个角落。"""
        img = cv.imread(file_path)  # 读取图片
        
        # 创建图像副本用于处理
        original = img.copy()
        img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)  # 将图片转换为灰度图
        
        # 添加高斯模糊减少噪声
        blurred = cv.GaussianBlur(img_gray, (5, 5), 0)
        
        # 使用Canny边缘检测代替简单二值化
        edges = cv.Canny(blurred, 75, 200)
        
        # 使用膨胀和腐蚀操作闭合边缘
        kernel = np.ones((5, 5), np.uint8)
        dilated = cv.dilate(edges, kernel, iterations=1)
        eroded = cv.erode(dilated, kernel, iterations=1)
        
        # 找到所有轮廓
        contours, hierarchy = cv.findContours(eroded, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        
        # 找到最大的周长的四边形的四个顶点
        max_perimeter = 0
        extreme_pnts = None
        
        # 按面积排序轮廓
        contours = sorted(contours, key=cv.contourArea, reverse=True)
        
        # 遍历轮廓寻找文档边界
        for contour in contours[:5]:  # 只检查最大的5个轮廓
            perimeter = cv.arcLength(contour, True)
            approx = cv.approxPolyDP(contour, 0.02 * perimeter, True)
            
            # 如果轮廓近似后有四个点，则认为是文档的四个角
            if len(approx) == 4:
                # 确保这是一个凸多边形(可能是文档)
                if cv.isContourConvex(approx) and perimeter > max_perimeter:
                    # 确保面积足够大(过滤掉小物体)
                    area = cv.contourArea(approx)
                    img_area = img.shape[0] * img.shape[1]
                    if area > 0.05 * img_area:  # 至少占据图像5%的面积
                        max_perimeter = perimeter
                        extreme_pnts = approx
        
        # 如果没有找到合适的四边形，使用整个图像的边界
        if extreme_pnts is None:
            h, w = img.shape[:2]
            extreme_pnts = np.array([
                [[0, 0]],
                [[w-1, 0]],
                [[w-1, h-1]],
                [[0, h-1]]
            ])
        
        corners = self.order_points(extreme_pnts.reshape(4, 2))  # 对四个顶点进行排序
        return img, corners  # 返回图片和排序后的顶点


    def order_points(self, pts):
        """对给定的四个点进行排序，返回排序后的点。"""
        rect = np.zeros((4, 2), dtype="float32")  # 初始化排序后的点为全零
        s = pts.sum(axis=1)  # 计算每个点的坐标和
        rect[0] = pts[np.argmin(s)]  # 左上角的点坐标和最小的那个点
        rect[2] = pts[np.argmax(s)]  # 右下角的点坐标和最大的那个点
        diff = np.diff(pts, axis=1)  # 计算每个点的坐标差
        rect[1] = pts[np.argmin(diff)]  # 右上角的点坐标差最小的那个点
        rect[3] = pts[np.argmax(diff)]  # 左下角的点坐标差最大的那个点

        print(f"0:左上角:{rect[0]},因为x,y的坐标和是: np.min(s)={np.min(s)}")
        print(f"1:右上角:{rect[1]},因为x,y的坐标差是: np.min(diff)={np.min(diff)}")
        print(f"2:右下角:{rect[2]},因为x,y的坐标和是: np.max(s)={np.max(s)}")
        print(f"3:左下角:{rect[3]},因为x,y的坐标差是: np.max(diff)={np.max(diff)}")

        return rect  # 返回排序后的点


    def crop_image(self, img, corners):
        """根据四个角落裁剪图片。"""

        top_left_corner = corners[0]
        top_right_corner = corners[1]
        bottom_right_corner = corners[2]
        bottom_left_corner = corners[3]

        # 计算目标图片的宽度和高度
        width, height = self.get_image_dimensions((top_left_corner, top_right_corner, bottom_right_corner, bottom_left_corner))  

        # 根据以上宽度和高度创建于原有四个点顺序对应的目标四个点的坐标
        dst_points = np.array([
            [0, 0],               # 左上角
            [width - 1, 0],       # 右上角
            [width - 1, height - 1], # 右下角
            [0, height - 1]       # 左下角
        ], dtype="float32")
            
        # 通过投影变换展平文档
        transform_matrix = cv.getPerspectiveTransform(corners, dst_points)
        dst = cv.warpPerspective(img, transform_matrix, (width, height))
        return dst  # 返回展平后的图像

    def get_image_dimensions(self, corners):
        """计算图片的宽度和高度，并确保适当的长宽比。"""
        top_left, top_right, bottom_right, bottom_left = corners
        
        # 计算宽度：取上边和下边的平均宽度
        width_top = np.sqrt(((top_right[0] - top_left[0]) ** 2) + ((top_right[1] - top_left[1]) ** 2))
        width_bottom = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + ((bottom_right[1] - bottom_left[1]) ** 2))
        width = max(int(width_top), int(width_bottom))
        
        # 计算高度：取左边和右边的平均高度
        height_left = np.sqrt(((bottom_left[0] - top_left[0]) ** 2) + ((bottom_left[1] - top_left[1]) ** 2))
        height_right = np.sqrt(((bottom_right[0] - top_right[0]) ** 2) + ((bottom_right[1] - top_right[1]) ** 2))
        height = max(int(height_left), int(height_right))
        
        # 标准A4纸长宽比约为1:1.414
        standard_ratio = 1.414
        current_ratio = height / width if width > 0 else 1
        
        # 如果当前比例与标准A4差距较大，则调整大小
        if abs(current_ratio - standard_ratio) > 0.5:
            # 按照检测到的宽度计算适当的高度
            if width > 0:
                height = int(width * standard_ratio)
        
        return width, height

    def preprocess_for_document(self, image):
        """对图像进行预处理以更好地识别文档"""
        # 转换为灰度图
        gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
        
        # 自适应阈值处理
        thresh = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                     cv.THRESH_BINARY, 11, 2)
        
        # 应用形态学操作增强文档边缘
        kernel = np.ones((5, 5), np.uint8)
        morph = cv.morphologyEx(thresh, cv.MORPH_CLOSE, kernel)
        
        # 应用双边滤波保留边缘信息
        filtered = cv.bilateralFilter(morph, 9, 75, 75)
        
        return filtered